{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "62549efe",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Abstract Syntax Tree\n",
    "\n",
    "`Kork` parses the code into the following AST https://github.com/langchain-ai/kork/blob/main/kork/ast.py.\n",
    "\n",
    "Below we'll show a few of the utilities to parse and print the ast."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "79fd8e41",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "nbsphinx": "hidden",
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": [
     "remove_cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "\n",
    "sys.path.insert(0, \"../\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fb92038-08ad-4cf6-91db-06fb25c281bd",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Parse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b5775536-ea67-4ed6-97e7-4d6f69f8d83a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kork.parser import parse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b6fec8a2-e4ce-48a8-b490-1e8be83886e4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Program(stmts=(Variable(name='x'),))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parse(\"x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4811e4e7-2a91-45f1-8d1f-3d9d75755698",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Program(stmts=(Assign(name='x', value=Literal(value=1)),))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parse(\"x = 1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b247fe8c-f782-4442-afb5-f84976aab87b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Program(stmts=(ExternFunctionDef(name='foo', params=ParamList(params=[]), return_type='Any', implementation=None, doc_string=''), Assign(name='x', value=Literal(value=1)), Assign(name='y', value=Binary(left=Binary(left=Variable(name='x'), operator='*', right=Variable(name='x')), operator='+', right=FunctionCall(name='foo', args=[])))))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parse(\n",
    "    \"\"\"\n",
    "extern fn foo() -> Any # Comment\n",
    "\n",
    "x = 1; // Comment\n",
    "y = x * x + foo()\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb6ac75e-6f3f-41b1-b0b8-c05db72a3549",
   "metadata": {},
   "source": [
    "## AstPrinter\n",
    "\n",
    "The AST printer can be used to translate AST syntax into a string representation of the program.\n",
    "\n",
    "This functionality is used to generate the prompt for the LLM as we'll see later in the docs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "62630c9b-eb0b-49a3-9bb8-f7bdfd3846b3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kork import AstPrinter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9f22fd0e-5390-4102-9141-00c1f76ab618",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "program = parse(\n",
    "    \"\"\"\n",
    "    extern fn foo() -> Any # Comment\n",
    "x=1; // Comment\n",
    "y=x*x + foo()\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9cb394cd-4d41-4ad0-944e-359ca70c4b40",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ParamList(params=[])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "program.stmts[0].params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "282a762d-5361-4d7d-b31d-b2f25912ff42",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "extern fn foo() -> Any\n",
      "x = 1\n",
      "y = x * x + foo()\n"
     ]
    }
   ],
   "source": [
    "print(AstPrinter().visit(program))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe62fa93-90f6-4554-82b4-72ba4f873158",
   "metadata": {},
   "source": [
    "## Nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76a1c4bf-34d7-4536-a08a-34cd7bac4901",
   "metadata": {},
   "source": [
    "The nodes in the ast are in the `ast` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "55051ae1-3490-4840-aed4-b8b54d7976d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kork.ast import FunctionCall, FunctionDef"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "37a7d6e7-5ea3-4c2d-8d3d-cd9e5c27a117",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mInit signature:\u001b[0m \u001b[0mFunctionCall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'str'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'Sequence[Expr]'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m      Represent a function call.\n",
       "\u001b[0;31mFile:\u001b[0m           ~/src/kork/kork/ast.py\n",
       "\u001b[0;31mType:\u001b[0m           ABCMeta\n",
       "\u001b[0;31mSubclasses:\u001b[0m     "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "?FunctionCall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "686ae81a-c866-4950-a852-f1288c7c6a9b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mInit signature:\u001b[0m\n",
       "\u001b[0mFunctionDef\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mname\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'str'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mparams\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'ParamList'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mbody\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'Sequence[Union[Stmt, Expr]]'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m    \u001b[0mreturn_type\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'str'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
       "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m      Represent a function definition with an implementation.\n",
       "\u001b[0;31mFile:\u001b[0m           ~/src/kork/kork/ast.py\n",
       "\u001b[0;31mType:\u001b[0m           ABCMeta\n",
       "\u001b[0;31mSubclasses:\u001b[0m     "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "?FunctionDef"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
